"""文本分类服务"""
import logging
from typing import Dict, Optional

from app.models.legalbert import classify_text
from app.schemas.requests import ClassificationResponse

logger = logging.getLogger(__name__)


async def classify_text_service(text: str, num_labels: int = 2) -> ClassificationResponse:
    """
    文本分类服务
    
    Args:
        text: 输入文本
        num_labels: 分类标签数量
        
    Returns:
        分类结果
    """
    try:
        # 执行分类
        predictions = classify_text(text, num_labels=num_labels)
        
        # 找到概率最高的标签
        predicted_label = max(predictions.items(), key=lambda x: x[1])[0]
        
        return ClassificationResponse(
            text=text,
            predictions=predictions,
            predicted_label=predicted_label
        )
    except Exception as e:
        logger.error(f"文本分类失败: {str(e)}", exc_info=True)
        raise ValueError(f"文本分类处理失败: {str(e)}")

